#!/bin/bash

set -euo pipefail

ASCEND_CANN_PACKAGE_PATH=/usr/local/Ascend/latest
export ASCEND_HOME_DIR=$ASCEND_CANN_PACKAGE_PATH

SHORT=r:,v:,p:,t,h:
LONG=run-mode:,soc-version:,perf:,test,hp:
OPTS=$(getopt -a --options $SHORT --longoptions $LONG -- "$@")
eval set -- "$OPTS"

TEST_MODE=0
HP_VALUE="1 1 256 128 256"

while :
do
    case "$1" in
        (-r | --run-mode )
            RUN_MODE="$2"
            shift 2;;
        (-v | --soc-version )
            SOC_VERSION="$2"
            shift 2;;
        (-p | --perf )
            IS_PERF="$2"
            shift 2;;
        (--test)
            TEST_MODE=1
            shift;;
        (-h | --hp)
            HP_VALUE="$2"
            shift 2;;
        (--)
            shift;
            break;;
        (*)
            echo "[ERROR] Unexpected option: $1";
            break;;
    esac
done


current_dir=$(dirname "$(realpath "\$0")")
echo "[INFO] current_dir directory is :$current_dir"

cd ../../ || { echo "[ERROR] Failed to change directory to ../../"; exit 1; }
act_dir=$(dirname "$(realpath "\$0")")
echo "act_dir directory is :$act_dir"

cd examples/09_lightning_attention/ || { echo "[ERROR] Failed to enter example/09_lightning_attention/"; exit 1; }

[ -f "a.out" ] && rm -f a.out

bishengcc "$current_dir/main.cpp" -arch $SOC_VERSION -I$act_dir \
    || { -d "[ERROR] bishengcc compilation failed."; exit 1; }

[ -d "build" ] && rm -rf build
[ -d "input" ] && rm -rf input
mkdir -p input
[ -d "output" ] && rm -rf output
mkdir -p output

mkdir -p bin
cd bin || { echo "[ERROR] Failed to enter bin directory"; exit 1; }

[ -f "../a.out" ] && mv ../a.out ./ascendc_matmul_bbit \
    || { echo "[ERROR] Failed to move a.out to ascendc_matmul_bbit"; exit 1; }

cd ../

if [ -n "$HP_VALUE" ]; then
    python3 scripts/gen_data.py --hp "$HP_VALUE" \
        || { echo "[ERROR] gen_data.py failed with $HP_VALUE", exit 1; }
else 
    python3 scripts/gen_data.py \
        || { echo "[ERROR] gen_data.py failed with default parameters", exit 1; }
fi

if [ -n "$HP_VALUE" ]; then
    ./bin/ascendc_matmul_bbit $HP_VALUE \
        || { echo "[ERROR] Excution failed with HP_VALUE: $HP_VALUE", exit 1; }
else
    ./bin/ascenc_matmul_bbit 1 1 256 128 256 \
        || { echo "[ERROR] Default execution failed"; exit 1; }
fi

if [ "$TEST_MODE" -eq 1 ]; then
    python3 -m scripts.cmp_data \
        || { echo "[ERROR] Test comparison failed"; exit 1;}
fi

echo "[SUCCESS] All steps completed successfully"